#!/usr/bin/env python3
import random
import sys
from scapy.all import *
from scapy.fields import *

# 源路由包头定义
class SourceRoute(Packet):
    fields_desc = [ 
        BitField("bos", 0, 1),    # 是否为最后一个条目
        BitField("port", 0, 15)   # 输出端口号
    ]

# 协议绑定
bind_layers(Ether, SourceRoute, type=0x1234)
bind_layers(SourceRoute, SourceRoute, bos=0)
bind_layers(SourceRoute, IP, bos=1)
bind_layers(SourceRoute, IPv6, bos=1)

def get_if():
    """获取eth0接口"""
    ifs = get_if_list()
    iface = next((i for i in ifs if "eth0" in i), None)
    if not iface:
        print("Cannot find eth0 interface")
        exit(1)
    return iface

def send_ipv4(src_ip, dst_ip, message, iface):
    """发送普通IPv4包"""
    pkt = Ether(src=get_if_hwaddr(iface), dst='ff:ff:ff:ff:ff:ff') / \
          IP(src=src_ip, dst=dst_ip) / \
          UDP(dport=4321, sport=random.randint(49152, 65535)) / \
          Raw(load=message.encode())
    print("Sending IPv4 packet:")
    pkt.show2()
    sendp(pkt, iface=iface, verbose=False)

def send_ipv6(src_ip, dst_ip, message, iface):
    """发送普通IPv6包（强制指定源地址）"""
    pkt = Ether(src=get_if_hwaddr(iface), dst='ff:ff:ff:ff:ff:ff') / \
          IPv6(src=src_ip, dst=dst_ip) / \
          UDP(dport=4321, sport=random.randint(49152, 65535)) / \
          Raw(load=message.encode())
    print("Sending IPv6 packet:")
    pkt.show2()
    sendp(pkt, iface=iface, verbose=False)

def send_source_route(src_ip, dst_ip, message, iface, ipv6=False):
    """发送源路由包（支持IPv4/IPv6）"""
    print("Enter source route path (space-separated port numbers)")
    print("Example for path S1->S3->S4: '3 2'")
    s = input("> ")
    
    try:
        ports = list(map(int, s.strip().split()))
        if not ports:
            raise ValueError
    except ValueError:
        print("Invalid port numbers")
        return

    # 构建源路由头
    pkt = Ether(src=get_if_hwaddr(iface), dst='ff:ff:ff:ff:ff:ff', type=0x1234)
    for i, port in enumerate(ports[:-1]):
        pkt = pkt / SourceRoute(bos=0, port=port)
    
    # 添加最后一跳和IP层
    if ipv6:
        pkt = pkt / SourceRoute(bos=1, port=ports[-1]) / \
              IPv6(src=src_ip, dst=dst_ip) / \
              UDP(dport=4321, sport=1234) / \
              Raw(load=message.encode())
    else:
        pkt = pkt / SourceRoute(bos=1, port=ports[-1]) / \
              IP(src=src_ip, dst=dst_ip) / \
              UDP(dport=4321, sport=1234) / \
              Raw(load=message.encode())

    print("Sending source routed packet:")
    pkt.show2()
    sendp(pkt, iface=iface, verbose=False)

def main():
    if len(sys.argv) < 5:
        print("Usage:")
        print("  IPv4:       send.py v4  <src_ip> <dst_ip> <message>")
        print("  IPv6:       send.py v6  <src_ipv6> <dst_ipv6> <message>")
        print("  SrcRoute:   send.py sr  <src_ip> <dst_ip> <message>")
        print("  SrcRouteV6: send.py sr6 <src_ipv6> <dst_ipv6> <message>")
        exit(1)

    mode = sys.argv[1]
    src_ip = sys.argv[2]
    dst_ip = sys.argv[3]
    message = sys.argv[4]
    iface = get_if()

    if mode == "v4":
        send_ipv4(src_ip, dst_ip, message, iface)
    elif mode == "v6":
        send_ipv6(src_ip, dst_ip, message, iface)
    elif mode == "sr":
        send_source_route(src_ip, dst_ip, message, iface)
    elif mode == "sr6":
        send_source_route(src_ip, dst_ip, message, iface, ipv6=True)
    else:
        print("Invalid mode")
        exit(1)

if __name__ == '__main__':
    main()